#ifndef _HBA_H
#define _HBA_H

#include <type.h>
#include <device/blkio.h>

#define AHCI_HBA_CLASS 0x10601
#define HBA_FIS_SIZE 256
#define HBA_CLB_SIZE 1024

/**
 * AHCI HBA Memory Registers
 * ref: SATA AHCI 3
*/
#define HBA_REG_CAP 0  // Host Capabilities
#define HBA_REG_GHC 1  // Generic Host Control
#define HBA_REG_IS 2   // Interrupt Status
#define HBA_REG_PI 3   // Ports Implemented
#define HBA_REG_VS 4   // Version

#define HBA_REG_PBASE (0x40)
#define HBA_REG_PSIZE (0x80 >> 2)

/**
 * AHCI HBA Port Registers (one set per port)
 * ref: SATA AHCI 3.3
*/
#define HBA_REG_PxCLB 0     // Command List Base Address
#define HBA_REG_PxCLBU 1    // Command List Base Address Upper 32 bits
#define HBA_REG_PxFB 2      // FIS Base Address
#define HBA_REG_PxFBU 3     // FIS Base Address Upper 32 bits
#define HBA_REG_PxIS 4      // Interrupt Status
#define HBA_REG_PxIE 5      // Interrupt Enable
#define HBA_REG_PxCMD 6     // Command and Status
#define HBA_REG_PxTFD 8     // Task File Data
#define HBA_REG_PxSIG 9     // Signature
#define HBA_REG_PxSSTS 10   // Serial ATA Status
#define HBA_REG_PxSCTL 11   // Serial ATA Control
#define HBA_REG_PxSERR 12   // Serial ATA Error
#define HBA_REG_PxSACT 13   // Serial ATA Active
#define HBA_REG_PxCI 14     // Command Issue
#define HBA_REG_PxSNTF 15   // Serial ATA Notification
#define HBA_REG_PxFBS 16    // FIS-Based Switch Control

#define HBA_PxCMD_FRE (1 << 4)
#define HBA_PxCMD_CR (1 << 15)
#define HBA_PxCMD_FR (1 << 14)
#define HBA_PxCMD_ST (1)

#define HBA_PxINTR_DMA (1 << 2)
#define HBA_PxINTR_DHR (1)
#define HBA_PxINTR_DPE (1 << 5)
#define HBA_PxINTR_TFEE (1 << 30)
#define HBA_PxINTR_IFE (1 << 27)

#define HBA_PxTFD_ERR (1)
#define HBA_PxTFD_BSY (1 << 7)
#define HBA_PxTFD_DRQ (1 << 3)

#define HBA_GHC_ACHI_ENABLE (1 << 31)
#define HBA_GHC_INTR_ENABLE (1 << 1)
#define HBA_GHC_RESET 1

#define HBA_PxSSTS_IPM(x) (((x) >> 8) & 0xF)
#define HBA_PxSSTS_SPD(x) (((x) >> 4) & 0xF)
#define HBA_PxSSTS_DET(x) ((x) & 0xF)

#define HBA_DEV_SIG_ATAPI 0xEB140101
#define HBA_DEV_SIG_ATA 0x00000101

#define HBA_CMDH_FIS_LEN(fis_bytes) (((fis_bytes) / 4) & 0x1F)
#define HBA_CMDH_ATAPI (1 << 5)
#define HBA_CMDH_WRITE (1 << 6)
#define HBA_CMDH_PREFETCH (1 << 7)
#define HBA_CMDH_R (1 << 8)
#define HBA_CMDH_CLR_BUSY (1 << 10)
#define HBA_CMDH_PRDT_LEN(entries) (((entries)&0xFFFF) << 16)

#define hba_clear_reg(reg) reg = -1

#define HBA_MAX_PRDTE 4

#define HBA_PxCMD_FRE (1 << 4)
#define HBA_PxCMD_CR (1 << 15)
#define HBA_PxCMD_FR (1 << 14)
#define HBA_PxCMD_ST (1)
#define HBA_PxINTR_DMA (1 << 2)
#define HBA_PxINTR_DHR (1)
#define HBA_PxINTR_DPS (1 << 5)
#define HBA_PxINTR_TFE (1 << 30)
#define HBA_PxINTR_HBF (1 << 29)
#define HBA_PxINTR_HBD (1 << 28)
#define HBA_PxINTR_IF (1 << 27)
#define HBA_PxINTR_NIF (1 << 26)
#define HBA_PxINTR_OF (1 << 24)
#define HBA_PxTFD_ERR (1)
#define HBA_PxTFD_BSY (1 << 7)
#define HBA_PxTFD_DRQ (1 << 3)

#define HBA_FATAL (HBA_PxINTR_TFE | HBA_PxINTR_HBF | HBA_PxINTR_HBD | HBA_PxINTR_IF)

typedef struct hba_command_header {
    u16 options;
    u16 prdt_len;
    u32 transferred_size;
    u32 cmd_table_base;
    u32 reserved[5];
} packed hba_cmdh_t;

#define HBA_PRDTE_BYTE_CNT(cnt) ((cnt & 0x3FFFFF) | 0x1)

typedef struct hba_physical_region_descriptor_table_entry {
    u32 data_base;
    u32 reserved[2];
    u32 byte_count;
} packed hba_prdte_t;

typedef struct hba_command_table {
    u8 command_fis[64];
    u8 atapi_cmd[16];
    u8 reserved[0x30];
    hba_prdte_t entries[3];
} packed hba_cmdt_t;

#define HBA_DEV_FEXTLBA 1
#define HBA_DEV_FATAPI (1 << 1)

struct hba_port;

typedef struct hba_device {
    char serial_num[20];
    char model[40];
    u32 flags;
    u64 max_lba;
    u32 block_size;
    u64 wwn;
    u8 cbd_size;
    struct {
        u8 error;
        u8 status;
    } last_result;
    u32 alignment_offset;
    u32 block_per_sec;
    u32 capabilities;
    struct hba_port* port;

    struct {
        int (*identify)(struct hba_device* dev);
        void (*submit)(struct hba_device* dev, blkio_req_t* req);
    } ops;
} hba_dev_t;

typedef struct hba_cmd_state {
    hba_cmdt_t* cmd_table;
    void* state_ctx;
} hba_cmds_t;

typedef struct hba_cmd_context {
    hba_cmds_t* issued[32];
    u32 tracked_ci;
} hba_cmd_ctx_t;

typedef struct hba_port {
    volatile u32* regs;
    u32 ssts;
    hba_cmdh_t* cmdlst;
    hba_cmd_ctx_t cmdctx;
    void* fis;
    hba_dev_t* device;
} hba_port_t;

typedef struct ahci_hba {
    volatile uptr* base;
    u32 ports_num;
    u32 ports_bmap;
    u32 cmd_slots;
    u32 version;
    hba_port_t* ports[32];
} hba_t;

int hba_init_device(hba_port_t* port);

int hba_prepare_cmd(hba_port_t* port, hba_cmdt_t** cmdt, hba_cmdh_t** cmdh);

void hba_reset_port(uptr* port_reg);

void hba_bind_buf(hba_cmdh_t* cmdh, hba_cmdt_t* cmdt, void* buf, size_t size);

void hba_bind_parted_buf(hba_cmdh_t* cmdh, hba_cmdt_t* cmdt, pbuffer_t* pbuf);

#endif /* _HBA_H */